#!/usr/bin/env python
"""
RNA design runner
"""

import os
from pyrna import utils
from pyrna import fsalib
from pyrna import ocamlfsa
from pyrna.automaton import DFA
from pyrna.parser import parseRNAStructure
from pyrna.param import Param

###############################################################################

class CfgRNADesign(object):
    """
    Design RNA under constraint
    """
    def __init__(self, param):
        assert isinstance(param, Param)
        self.param = param
        self.out_fsa = self.param.outprefix + ".fsa"
        self.out_fsa_seq = self.out_fsa+"_%d-%d.seq" % (self.param.len, self.param.len)
        self.out_grammar = self.param.outprefix + ".ggd"
        self.out_grammar_seq = self.out_grammar + ".seq"
        
        self.fsa_report = {'start':0, 'final':0, 'state':0, 'transition':0}
        self.has_base_constraint = False
        self.is_out_fsa_empty = True
        self.is_grammar_empty = False
        
        self.e_struct = self.param.struct
        self.e_seq = self.param.seq
        self.mda_motif = self.param.mda_motif
        self.file_fbdn_motifs = self.param.fbdn
        
        self.tick = utils.Tick()
        self.num_seq_fsa = None
        self.num_seq_grammar = None
        
    def __str__(self):
        return """\
#seq_fsa: %s, #seq_grammar: %s
""" % (list(map(str, [self.num_seq_fsa, self.num_seq_grammar])))
        
    def design(self):
        """design sequences"""
        self._make_fsa()
        self._plot_fsa()
        self._parse_report()
        if self.param.fsa_seq:
            self.num_seq_fsa = self._generate_fsa_seq()
        self._make_grammar()
        self.num_seq_grammar = self._generate_grammar_seq()    
    
    def get_num_candidates(self):
        """Return number of all possible sequences"""
        return self.num_seq_grammar
    
    def get_fsa_seq(self):
        return self.out_fsa_seq
    
    def get_grammar_seq(self):
        return self.out_grammar_seq
    
    def _make_fsa(self):
        """Generate FSA for motif constraint"""
        self.has_base_constraint = fsalib.has_base_constraint(self.e_seq)
        utils.info('Expected sequence has base constraint?: %s' % self.has_base_constraint)
        utils.info('file_fbdn_motifs %s, mda_motif: %s'
            % (self.file_fbdn_motifs, self.mda_motif))
        
        if self.file_fbdn_motifs and self.mda_motif:
            utils.info('FSA with forbidden motifs and mandatory motif ...')
            fsalib.designSeq(self.mda_motif, self.e_seq, self.file_fbdn_motifs, 
                self.out_fsa)
            
        elif self.file_fbdn_motifs and self.has_base_constraint:
            utils.info('FSA with forbidden motifs and base constraints ...')
            fsa_base = self.out_fsa+'.base'
            fsalib.seqToFsa(self.e_seq, fsa_base)
            fsa_fbdn = fsalib.genFsaNotForbiddenMotifs(self.file_fbdn_motifs)
            ocamlfsa.inter(fsa_base, fsa_fbdn, self.out_fsa)    
                
        elif self.file_fbdn_motifs:
            utils.info('FSA with forbidden motifs ...')
            self.out_fsa = fsalib.genFsaNotForbiddenMotifs(self.file_fbdn_motifs)
        else:
            utils.info('FSA without forbidden motifs ...')
            fsalib.seqToFsa(self.e_seq, self.out_fsa)
        
        self.is_out_fsa_empty = fsalib.isEmptyFsa(self.out_fsa)
        utils.info(self.tick.tock())
    
    def _plot_fsa(self):
        """Plot FSA"""
        if self.param.plot:
            ocamlfsa.plotFsa(self.out_fsa, 'png')
    
    def _parse_report(self):
        """Parse FSA summary"""
        num_start, num_final, num_states, num_trans = ocamlfsa.report(self.out_fsa)
        utils.info('FSA report: start=%d, final=%d, states=%d, trans=%d'
            % (num_start, num_final, num_states, num_trans))
        self.fsa_report['start'] = num_start
        self.fsa_report['final'] = num_final
        self.fsa_report['state'] = num_states
        self.fsa_report['transition'] = num_trans
        
    def _generate_fsa_seq(self):
        """Generate sequences from FSA"""
        num_seq_fsa = None    
        if self.is_out_fsa_empty:
            utils.info('Empty FSA, no sequence satisfying the constraints.')
            num_seq_fsa = 0
        elif self.param.fsa_seq: # generate if necessary
            ocamlfsa.generate(self.out_fsa, self.param.len, self.param.len, 
                self.out_fsa_seq)
            num_seq_fsa = utils.wc_line(self.out_fsa_seq)
            utils.info('%d sequences written to %s' % (
                num_seq_fsa, self.out_fsa_seq))
        
        utils.info(self.tick.tock())
        return num_seq_fsa
    
    def _make_grammar(self):
        """Generate Grammar for structure constraint"""
        struct_g = parseRNAStructure(self.e_struct)
        utils.debug(struct_g)
    
        if not self.is_out_fsa_empty:
            dfa = DFA("ACGU") 
            dfa.readFsa(self.out_fsa)
    
            # do intersection
            inter_g = struct_g.intersect(dfa)
            utils.info('#NT=%d, #PROD=%d' % (inter_g.num_nt(), inter_g.num_prod()))
            utils.info('Grammar intersect FSA finished: time=%s' % self.tick.tock())
            utils.debug(inter_g)
            
            inter_g.minimize()
            utils.info('#NT=%d, #PROD=%d' % (inter_g.num_nt(), inter_g.num_prod()))
            utils.info('Intersected grammar minimized: time=%s' % self.tick.tock())
            
            if inter_g.num_prod() != 0:
                inter_g.toGenRGenS(self.out_grammar)
            else:
                utils.info('Empty grammar.')
                self.is_grammar_empty = True
                os.system("touch %s" % self.out_grammar)
                
        elif not (self.has_base_constraint or self.file_fbdn_motifs or self.mda_motif):
            utils.info('make grammar for structure only.')
            struct_g.toGenRGenS(self.out_grammar)
            
        utils.info(self.tick.tock())
        
    def _generate_grammar_seq(self):
        """Generate sequences by GrGFreqs"""
        if self.is_grammar_empty:
            os.system("touch %s" % self.out_grammar_seq)
            num_seq_grammar = 0
        else:
            len_struct = len(self.e_struct)
            command = fsalib.GrGFreqs_COMMAND % (self.out_grammar, 
                len_struct, self.out_grammar_seq)
            
            if self.param.gen is not None:
                command = "grgfreqs --draw %s %d %d -o %s 2>/dev/null" % (
                    self.out_grammar, len_struct, self.param.gen, self.out_grammar_seq)
            
            utils.info(command)
            os.system(command)
            num_seq_grammar = utils.wc_line(self.out_grammar_seq)
            
        utils.info('End of generation')
        utils.info('%d sequences for grammar written to %s'
                   % (num_seq_grammar, self.out_grammar_seq))
        utils.info(self.tick.tock())
    
        return num_seq_grammar
    
###############################################################################
 